package com.streaming.flink.udtagg;

import org.apache.commons.compress.utils.Lists;
import org.apache.flink.calcite.shaded.com.google.common.collect.Maps;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.table.api.*;
import org.apache.flink.types.Row;
import org.apache.flink.types.RowKind;
import org.apache.flink.util.CloseableIterator;

import java.time.LocalDate;
import java.util.List;
import java.util.Map;
import java.util.function.BiFunction;

import static org.apache.flink.table.api.Expressions.$;
import static org.apache.flink.table.api.Expressions.call;

/**
 * UDTAGG函数目前只有在流模式下使用
 */
public class UDTAGGFunctionExample {
    public static void main(String[] args) throws Exception {
        // setup the environment
        final EnvironmentSettings settings = EnvironmentSettings.newInstance()
                .useBlinkPlanner()
                .inStreamingMode()
                .build();
        final TableEnvironment tEnv = TableEnvironment.create(settings);
        Configuration configuration = tEnv.getConfig().getConfiguration();
        configuration.setString("table.exec.mini-batch.enabled", "true");
        configuration.setString("table.exec.mini-batch.allow-latency", "5 s");
        configuration.setString("table.exec.mini-batch.size", "5000");
        configuration.setString("sql-client.execution.result-mode", "table");

        // execute different kinds of functions
        executeUDTAGGFunction(tEnv);
    }

    private static void executeUDTAGGFunction(TableEnvironment tEnv) {
        // create a table with example data
        final Table customers =
                tEnv.fromValues(
                        DataTypes.of("ROW<name STRING, order_date DATE, item_count INT,item_price INT>"),
                        Row.of("Guillermo Smith", LocalDate.parse("2020-12-01"), 3, 2),
                        Row.of("Guillermo Smith", LocalDate.parse("2020-12-05"), 5, 3),
                        Row.of("Guillermo Smith", LocalDate.parse("2020-03-23"), 4, 3),
                        Row.of("Guillermo Smith", LocalDate.parse("2020-06-02"), 10, 2),
                        Row.of("Leann Holloway", LocalDate.parse("2020-05-26"), 9, 1),
                        Row.of("Leann Holloway", LocalDate.parse("2020-05-27"), 1, 2),
                        Row.of("Leann Holloway", LocalDate.parse("2020-10-14"), 1, 3),
                        Row.of("Leann Holloway", LocalDate.parse("2020-10-02"), 12, 3),
                        Row.of("Ellen Ortega", LocalDate.parse("2020-06-18"), 100, 3));
        tEnv.createTemporaryView("customers", customers);


        // register and execute the function
        tEnv.createTemporarySystemFunction("top2", Top2Function.class);
//        tEnv.from("customers")
//                .groupBy($("name"))
//                .flatAggregate(call("top2", $("item_price")).as("item_price", "rank"))
//                .select($("name"), $("item_price"), $("rank"))
//                .execute().print();

        Table topRes = tEnv.from("customers")
                .groupBy($("name"))
                .flatAggregate(call(Top2Function.class, $("item_price")).as("item_price", "rank"))
                .select($("name"), $("item_price"), $("rank"));

        tEnv.createTemporaryView("topRes", topRes);
        topRes.printSchema();

        TableResult res = tEnv.executeSql("select * from topRes");
        CloseableIterator<Row> it = res.collect();
        Map<String, List<Row>> map = Maps.newHashMap();
        while (it.hasNext()) {
            Row row = it.next();
            if(row.getKind()== RowKind.INSERT){
                String name = row.getField("name").toString();
                map.compute(name, new BiFunction<String, List<Row>, List<Row>>() {
                    @Override
                    public List<Row> apply(String s, List<Row> rows) {
                        if(rows == null){
                            rows = Lists.newArrayList();
                            rows.add(row);
                            rows.add(null);
                        }else {
                            rows.set(1,rows.get(0));
                            rows.set(0,row);
                        }
                        return rows;
                    }
                });
            }
        }
        for(Map.Entry<String,List<Row>> row:map.entrySet()){
            System.out.println(String.format("%s,%s\n",row.getKey(),row.getValue()));
        }

        // clean up
        tEnv.dropTemporaryView("topRes");
        // clean up
        tEnv.dropTemporaryView("customers");
    }
}
